import sys
from PyQt5.Qt import *
from queue import Queue
import requests
import threading

#设置定时器,将过滤的url结果存放到PlainTextEdit中
U = Queue()
class MyObject(QObject):
    def timerEvent(self, evt):
        global obj
        global time_id
        global progressbar
        global consume_count
        #flag是一个标志位,用于判断一次扫描是否结束
        global flag
        if flag == 1:
            progressbar.reset()
            consume_count = 0
        progressbar.setValue(consume_count)
        global U
        global pte
        a = U.empty()
        if a == False:
            pte.appendHtml(U.get())
            U.task_done()
'''重写类方便以后扩展'''
class Window(QWidget):
    pass
class PushButton(QPushButton):
    pass
class Label(QLabel):
    pass
class LineEdit(QLineEdit):
    pass
class Frame(QFrame):
    pass
class ComboBox(QComboBox):
    pass
class CheckBox(QCheckBox):
    def __init__(self, parent):
        super().__init__(parent)
        #取消QCheckBox的三态
        self.setTristate(False)
class ProgressBar(QProgressBar):
    pass
class PlainTextEdit(QPlainTextEdit):
    pass
class MessageBox(QMessageBox):
    pass

#扫描线程
consume_count = 0
dir_count = 0
Q = Queue()
domain = ""
status_code_num = 0
timeout = 0
flag = 0
thread_num = 0
time_id = None
lock = threading.Lock()
def task():
    global consume_count
    global status_code_num
    global domain
    global timeout
    global dir_count
    global Q
    global UA
    global P
    global p_count
    global ua_count
    global U
    global lock
    global progressbar
    global flag
    while dir_count:
        if Q.empty():
            break
        if flag == 1:
            break
        lock.acquire()
        consume_count += 1
        lock.release()
        dir_count -= 1
        timeout = int(timeout)
        url = domain.strip() + Q.get().replace('\n', '')
        Q.task_done()
        s = requests.session()
        s.keep_alive = False
        html = None
        try:
            if UA == []:
                headers = {'User-Agent': 'Mozilla/5.0 (X11; Linux x86_64; rv:68.0) Gecko/20100101 Firefox/68.0'}
                if P == []:
                    html = s.get(url, headers=headers, timeout=timeout)
                else:
                    p = range(0, p_count)
                    P[p] = P[p].replace('\n', '')
                    proxies = {'https': 'https://' + P[p], 'http': 'http://' + P[p]}
                    html = s.get(url, headers=headers, proxies=proxies, timeout=timeout)
            else:
                u = range(0, ua_count)
                ua = UA[u].replace('\n', '')
                headers = {'User-Agent': ua}
                if P == []:
                    html = s.get(url, headers=headers, timeout=timeout)
                else:
                    p = range(0, p_count)
                    P[p] = P[p].replace('\n', '')
                    proxies = {'https': 'https://' + P[p], 'http': 'http://' + P[p]}
                    html = s.get(url, headers=headers, proxies=proxies, timeout=timeout)
        except:
            continue
        if status_code_num == 1:
            if html.status_code == 200:
                U.put("<a>" + url + "</a>" + "     200")
        if status_code_num == 2:
            if html.status_code >= 300 and html.status_code <= 399:
                U.put("<a>" + url + "</a>" + "      " + str(html.status_code))
        if status_code_num == 4:
            if html.status_code == 403:
                U.put("<a>" + url + "</a>" + "      403")
        if status_code_num == 3:
            if html.status_code == 200:
                U.put("<a>" + url + "</a>" + "     200")
            if html.status_code >= 300 and html.status_code <= 399:
                U.put("<a>" + url + "</a>" + "      " + str(html.status_code))
        if status_code_num == 5:
            if html.status_code == 200:
                U.put("<a>" + url + "</a>" + "     200")
            if html.status_code == 403:
                U.put("<a>" + url + "</a>" + "      403")
        if status_code_num == 6:
            if html.status_code >= 300 and html.status_code <= 399:
                U.put("<a>" + url + "</a>" + "      " + str(html.status_code))
            if html.status_code == 403:
                U.put("<a>" + url + "</a>" + "      403")
        if status_code_num == 7:
            if html.status_code == 200:
                U.put("<a>" + url + "</a>" + "     200")
            if html.status_code >= 300 and html.status_code <= 399:
                U.put("<a>" + url + "</a>" + "      " + str(html.status_code))
            if html.status_code == 403:
                U.put("403        "+"<a href="+url+">" + url + "</a>")



flag = 0
#退出按钮触发事件
def exit_scan():
    global flag
    global start_button
    global shut_button
    global info_label
    global progressbar
    progressbar.setValue(0)
    info_label.setText("已停止本次扫描")
    flag = 1
    start_button.setEnabled(True)
    shut_button.setEnabled(False)
    proxy_button.setEnabled(True)
    ua_button.setEnabled(True)

def proxy_thread():
    t = threading.Thread(target=proxy_import)
    t.start()
#使用生产者消费者模型,实现对同一字典的多线程安全读写
proxy_filename = ""
p_count = 0
P = []
def proxy_import():
    global P
    global proxy_filename
    global p_count
    proxy_filename = QFileDialog.getOpenFileName()[0]
    if proxy_filename != "":
        info_label.setText("正在导入代理....")
        with open(proxy_filename, "r") as f:
            while True:
                data = f.readline()
                if data:
                    P.append(data)
                    p_count += 1
                else:
                    break
        info_label.setText("导入代理完成")
def ua_thread():
    t = threading.Thread(target=ua_import)
    t.start()
ua_filename = ""
ua_count = 0
UA = []
def ua_import():
    global ua_filename
    ua_filename = QFileDialog.getOpenFileName()[0]
    global ua_count
    global UA
    if ua_filename != "":
        info_label.setText("正在导入UA....")
        with open(ua_filename, "r") as f:
            while True:
                data = f.readline()
                if data:
                    UA.append(data)
                    ua_count += 1
                else:
                    break
        info_label.setText("导入UA完成")
Q = Queue()
def start_scan():
    #检测是否输入域名及选择按钮
    # 获取域名
    global flag
    global domain
    domain = domain_edit.text()
    global status_code_num
    global timeout
    global thread_num
    global progressbar
    global shut_button
    global pte
    proxy_button.setEnabled(False)
    ua_button.setEnabled(False)
    status_code_num = 0
    pte.clear()
    flag = 0
    # 获取字典类型和状态码
    shut_button.setEnabled(True)
    scan_type = []
    for checkbox in window.children():
        if isinstance(checkbox, CheckBox):
            if checkbox.objectName() == "status_code":
                if checkbox.checkState() == 2:
                    status_code_num += checkbox.property("value")
            else:
                if checkbox.checkState() == 2:
                    if checkbox.property("id") == 1:
                        scan_type.append("DIR")
                    if checkbox.property("id") == 2:
                        scan_type.append("PHP")
                    if checkbox.property("id") == 3:
                        scan_type.append("ASP")
                    if checkbox.property("id") == 4:
                        scan_type.append("MDB")
                    if checkbox.property("id") == 5:
                        scan_type.append("ASPX")
                    if checkbox.property("id") == 6:
                        scan_type.append("JSP")
    if domain == "":
        mb = MessageBox(window)
        mb.setWindowTitle("警告")
        mb.setText("^_^请输入即将被扫描的域名")
        mb.show()
    else:
        if scan_type == []:
            mb = MessageBox(window)
            mb.setWindowTitle("警告")
            mb.setText("^_^请选择扫描的字典类型")
            mb.show()
        else:
            if status_code_num == 0:
                mb = MessageBox(window)
                mb.setWindowTitle("警告")
                mb.setText("^_^请选择探测码类型")
                mb.show()
    thread_num = thread_combobox.currentText()
    timeout = timeout_combobox.currentText()
    #介于字典一般不会大于内存, 采用一次性装入队列
    #使用dir_count获取字典总条数
    global Q
    Q = Queue()
    global dir_count
    for i in scan_type:
        with open("./config/" + i + ".txt", "r") as f:
            while True:
                data = f.readline()
                if data:
                    Q.put(data)
                    dir_count += 1
                else:
                    info_label.setText("正在装载字典,请稍等....")
                    break
    thread_num = int(thread_num)
    if domain == "" or status_code_num == 0 or scan_type == []:
        pass
    else:
        progressbar.setMinimum(0)
        progressbar.setMaximum(dir_count)
        start_button.setEnabled(False)
        info_label.setText("开始扫描")
        for i in range(0,thread_num):
            t = threading.Thread(target=task,args=())
            #t.setDaemon(True)
            t.start()


# 程序入口
app = QApplication(sys.argv)
window = Window()
# 作者信息
window.setWindowTitle("山一后台扫描工具  By:Corn WX:a2635918064")
# 图标
icon = QIcon("./logo.png")
window.setWindowIcon(icon)
# 设置主窗体的固定大小
window.setFixedSize(800, 600)
# 透明度设置
window.setWindowOpacity(0.95)
# 开始扫描按钮
start_button = PushButton(window)
start_button.setText("开始扫描")
start_button.resize(90, 40)
start_button.move(560, 22)
start_button.pressed.connect(start_scan)
# 结束扫描按钮
shut_button = PushButton(window)
shut_button.setText("停止扫描")
shut_button.resize(90, 40)
shut_button.move(560, 82)
shut_button.pressed.connect(exit_scan)
shut_button.setEnabled(False)
# 导入代理按钮
proxy_button = PushButton(window)
proxy_button.resize(90, 40)
proxy_button.setText("导入代理")
proxy_button.move(680, 82)
proxy_button.pressed.connect(proxy_thread)
# 导入UA按钮
ua_button = PushButton(window)
ua_button.setText("导入UA")
ua_button.resize(90, 40)
ua_button.move(680, 22)
ua_button.pressed.connect(ua_thread)
# 提示信息词
info_label = Label(window)
info_label.setText("提示信息:企鹅后台扫描工具准备就绪...")
info_label.setStyleSheet("font-size:12px;")
info_label.move(30, 145)
# 域名输入框提示词
domain_label = Label(window)
domain_label.setText("域名:")
domain_label.move(30, 35)
domain_label.setStyleSheet("font-size:14px;")
# 域名输入框
domain_edit = LineEdit(window)
domain_edit.resize(450, 25)
domain_edit.move(75, 30)
# 线程数选拉框提示词
thread_label = Label(window)
thread_label.setText("线程数:")
thread_label.setStyleSheet("font-size:15px;")
thread_label.move(30, 65)
# 线程数选拉框
thread_combobox = ComboBox(window)
thread_combobox.resize(40, 20)
thread_combobox.move(90, 63)
for i in range(1, 21):
    thread_combobox.addItem(str(i))
# 超时选拉框提示词
timeout_label = Label(window)
timeout_label.setText("超时时间:")
timeout_label.setStyleSheet("font-size:15px;")
timeout_label.move(30, 95)
# 超时选拉框
timeout_combobox = ComboBox(window)
timeout_combobox.resize(40, 20)
timeout_combobox.move(105, 93)
for i in range(1, 21):
    timeout_combobox.addItem(str(i))
# 扫描类型选项
dir_checkbox = CheckBox(window)
dir_checkbox.setText("DIR")
dir_checkbox.move(220, 65)
php_checkbox = CheckBox(window)
php_checkbox.setText("PHP")
php_checkbox.move(220, 85)
asp_checkbox = CheckBox(window)
asp_checkbox.setText("ASP")
asp_checkbox.move(220, 105)
mdb_checkbox = CheckBox(window)
mdb_checkbox.setText("MDB")
mdb_checkbox.move(280, 65)
aspx_checkbox = CheckBox(window)
aspx_checkbox.setText("ASPX")
aspx_checkbox.move(280, 85)
jsp_checkbox = CheckBox(window)
jsp_checkbox.setText("JSP")
jsp_checkbox.move(280, 105)
dir_checkbox.setProperty("id", 1)
php_checkbox.setProperty("id", 2)
asp_checkbox.setProperty("id", 3)
mdb_checkbox.setProperty("id", 4)
aspx_checkbox.setProperty("id", 5)
jsp_checkbox.setProperty("id", 6)
# 探测类型
status_200 = CheckBox(window)
status_200.setText("探测200")
status_200.move(410, 65)
status_403 = CheckBox(window)
status_403.setText("探测403")
status_403.move(410, 85)
status_3xx = CheckBox(window)
status_3xx.setText("探测3XX")
status_3xx.move(410, 105)
status_200.setObjectName("status_code")
status_3xx.setObjectName("status_code")
status_403.setObjectName("status_code")
#与Linux中的chmod权限值同理
status_200.setProperty("value", 1)
status_3xx.setProperty("value", 2)
status_403.setProperty("value", 4)
# 扫描结果
pte = PlainTextEdit(window)
pte.resize(730, 400)
pte.move(30, 166)
pte.setReadOnly(True)
'''
tcf = QTextCharFormat()
tcf.setFontUnderline(True)
pte.setCurrentCharFormat(tcf)
'''
# 进度条
progressbar = ProgressBar(window)
progressbar.resize(760, 20)
progressbar.move(30, 570)
obj = MyObject()
time_id = obj.startTimer(100)
window.show()
sys.exit(app.exec_())
